from solvers.solverbase import SolverBase
from solvers.analytical.analyticalbase import AnalyticalBase
from dolfin import *
from math import sqrt
import numpy

class GaussLegendre4(SolverBase,AnalyticalBase):
    '2-symplectic method rewritten without increments.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        AnalyticalBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        p = self.p
        V = FunctionSpace(mesh,'CG',p)
        VV = VectorFunctionSpace(mesh,'CG',p)

        phi_ = interpolate(problem.phi_,V)
        phiP = Function(V)                  # solution to predictor system
        phiC = Function(V)                  # solution to corrector system
        phi = TrialFunction(V)          
        varphi = TestFunction(V)

        t = problem.t          
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)

        (ib,ib_value,bParts) = problem.get_inflowBCS()
        if problem.weakBC:
            aP = phi*varphi*dx
            lP = phi_*varphi*dx - dt*sqrt(3.)/6.*(inner(nabla_grad(varphi),phiC*v)*dx + phiC*varphi*nabla_div(v)*dx)\
                                + dt*sqrt(3.)/6.*(inner(nabla_grad(varphi),phi_*v)*dx + phi_*varphi*nabla_div(v)*dx)

            aC = phi*varphi*dx - 0.5*dt*(inner(nabla_grad(varphi),phi*v)*dx + phi*varphi*nabla_div(v)*dx)
            lC = phi_*varphi*dx + dt*(0.5-sqrt(3.)/6.)*(inner(nabla_grad(varphi),phi_*v)*dx + phi_*varphi*nabla_div(v)*dx)\
                                + dt*(sqrt(3.)/6.)*(inner(nabla_grad(varphi),phiP*v)*dx + phiP*varphi*nabla_div(v)*dx)
        else:
            n = FacetNormal(mesh)
            aP = phi*varphi*dx #- dt*sqrt(3.)/12.*(phi*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi*v)*dx - phi*varphi*nabla_div(v)*dx)
            lP = phi_*varphi*dx + dt*sqrt(3.)/6.*(phiC*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phiC*v)*dx - phiC*varphi*nabla_div(v)*dx)\
                                - dt*sqrt(3.)/6.*(phi_*varphi*inner(v,n)*ds - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)

            aC = phi*varphi*dx + 0.5*dt*(phi*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi*v)*dx - phi*varphi*nabla_div(v)*dx)
            lC = phi_*varphi*dx -dt*(0.5-sqrt(3.)/6.)*(phi_*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)\
                                -dt*(sqrt(3.)/6.)*(phiP*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phiP*v)*dx - phiP*varphi*nabla_div(v)*dx)

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0

        phiC.assign(phi_)
        
        phiC_ = Function(V)
        phiC_.assign(phi_)

        AP = assemble(aP)
        AC = assemble(aC,exterior_facet_domains=bParts)
        phi = Function(V)
   
        t = float(dt)
        count = 1
       
        iteration = 0
        maxIter = 10
        tol = 1E-8
        eps = 1.
        
        while t <= T:
            ibc = DirichletBC(V,ib_value(t),ib)
            while eps > tol and iteration < maxIter:
                iteration += 1
                LP = assemble(lP)
                solve(AP,phi.vector(),LP)  # solve system and assign the solution to phiP
                phiP.assign(phi)

                LC = assemble(lC,exterior_facet_domains=bParts)
                ibc.apply(AC,LC)
                solve(AC,phi.vector(),LC)  # solve system and assign the solution to phiC
                phiC.assign(phi)

                diff = phiC.vector().array() - phiC_.vector().array()
                eps = numpy.linalg.norm(diff,ord=2)        # compute the error
                phiC_.assign(phiC)
                
        #        print iteration, eps

            phi_.assign(phiC)
            self.update(problem,phi_,t,float(dt),count)

            t += float(dt)
            count += 1
            eps = 1.
            iteration = 0

        problem.save_data(self.saveDir,str(self.CFL))

#<----------------------------------------------TRICKY THINGS---------------------------------- 
class GaussLegendre4Increment(SolverBase,AnalyticalBase):
    '2-symplectic method as found on wiki. Does not work well with boundary conditions! '
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        AnalyticalBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        p = self.p
        V = FunctionSpace(mesh,'CG',p)
        VV = VectorFunctionSpace(mesh,'CG',p)

        phi_ = interpolate(problem.phi_,V)
        k1 = Function(V)                   
        k2 = Function(V)                   
        phi = TrialFunction(V)          
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)

        (ib,ib_value,bParts) = problem.get_inflowBCS()
        if problem.weakBC:
            a = phi*varphi*dx - 0.25*dt*(inner(nabla_grad(varphi),phi*v)*dx + phi*varphi*nabla_div(v)*dx)
            l1 =  + dt*(inner(nabla_grad(varphi),phi_*v)*dx + phi_*varphi*nabla_div(v)*dx)\
                  + dt*(0.25-sqrt(3.)/6.)*(inner(nabla_grad(varphi),k2*v)*dx + k2*varphi*nabla_div(v)*dx)

            l2 =  + dt*(inner(nabla_grad(varphi),phi_*v)*dx + phi_*varphi*nabla_div(v)*dx)\
                  + dt*(0.25+sqrt(3.)/6.)*(inner(nabla_grad(varphi),k1*v)*dx + k1*varphi*nabla_div(v)*dx)
        else:
            n = FacetNormal(mesh)
            a = phi*varphi*dx + 0.25*dt*(phi*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi*v)*dx - phi*varphi*nabla_div(v)*dx)
            l1 =  - dt*(phi_*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)\
                  - dt*(0.25-sqrt(3.)/6.)*(k2*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),k2*v)*dx - k2*varphi*nabla_div(v)*dx)

            l2 =  - dt*(phi_*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx)\
                  - dt*(0.25+sqrt(3.)/6.)*(k1*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),k1*v)*dx - k1*varphi*nabla_div(v)*dx)

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0

        k2.assign(interpolate(Constant(0),V)) # intial guess for the value
        
        k1_ = Function(V)
        k2_ = Function(V) # to hold solutions in itarating
        k1_.assign(k2)
        k2_.assign(k2)

        A = assemble(a,exterior_facet_domains=bParts)
        phi = Function(V)
   
        t = float(dt)
        count = 1
        
        iteration = 0
        maxIter = 10
        tol = 1E-8
        eps = 1.

        while t <= T:
            ibc = DirichletBC(V,Constant(0.),ib)
            while eps > tol and iteration < maxIter:
                iteration += 1
                L1 = assemble(l1,exterior_facet_domains=bParts)
                ibc.apply(L1)
                solve(A,phi.vector(),L1)
                k1.assign(phi)
                diff1 = k1.vector().array() - k1_.vector().array() 
                eps1 = numpy.linalg.norm(diff1,ord=2) 

                L2 = assemble(l2,exterior_facet_domains=bParts)
                ibc.apply(L2)
                solve(A,phi.vector(),L2)
                k2.assign(phi)
                diff2 = k2.vector().array() - k2_.vector().array() 
                eps2 = numpy.linalg.norm(diff2,ord=2) 

                eps = numpy.amax([eps1,eps2])
                k1_.assign(k1)
                k2_.assign(k2)
                print iteration, eps, t
            
            ibc = DirichletBC(V,ib_value(t),ib)
            ibc.apply(phi_.vector())
            phi_.vector()[:] = phi_.vector().array() + 0.5*k1.vector().array() + 0.5*k2.vector().array()
            
            self.update(problem,phi_,t,float(dt),count)

            t += float(dt)
            count += 1
            eps = 1.
            iteration = 0

        problem.save_data(self.saveDir,str(self.CFL))
